package org.loong.crypto.extension.cms.jce.io;

import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.Key;
import java.security.NoSuchAlgorithmException;

import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.crypto.io.CipherIOException;
import org.bouncycastle.crypto.io.InvalidCipherTextIOException;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.algorithm.CipherAlgorithm;
import org.loong.crypto.core.params.CipherParameters;
import org.loong.crypto.extension.cms.EnvelopedDataHelper;
import org.loong.crypto.service.CryptoService;

public class CipherInputStream extends FilterInputStream {

    private final EnvelopedDataHelper helper = new EnvelopedDataHelper();

    private final CryptoService cryptoService;

    private final ASN1ObjectIdentifier contentEncryptionAlgorithm;

    private final Key secretKey;

    private final byte[] iv;

    private final byte[] inputBuffer = new byte[8192];

    private boolean finalized = false;

    private byte[] buf;

    private int maxBuf;

    private int bufOff;

    /**
     * Constructs a CipherInputStream from an InputStream and an initialised Cipher.
     * 
     * @param input the input
     * @param cryptoService the cryptoService
     * @param contentEncryptionAlgorithm the contentEncryptionAlgorithm
     * @param secretKey the secretKey
     * @param iv the iv
     */
    public CipherInputStream(InputStream input, CryptoService cryptoService, ASN1ObjectIdentifier contentEncryptionAlgorithm, Key secretKey, byte[] iv) {
        super(input);

        this.cryptoService = cryptoService;
        this.contentEncryptionAlgorithm = contentEncryptionAlgorithm;
        this.secretKey = secretKey;
        this.iv = iv;
    }

    /**
     * Read data from underlying stream and process with cipher until end of stream or some data is available after
     * cipher processing.
     *
     * @return -1 to indicate end of stream, or the number of bytes (> 0) available.
     */
    private int nextChunk() throws IOException {
        if (finalized) {
            return -1;
        }

        bufOff = 0;
        maxBuf = 0;

        // Keep reading until EOF or cipher processing produces data
        while (maxBuf == 0) {
            int read = in.read(inputBuffer);
            if (read == -1) {
                buf = finaliseCipher();
                if ((buf == null) || (buf.length == 0)) {
                    return -1;
                }
                maxBuf = buf.length;
                return maxBuf;
            }

            try {
                byte[] buffer = new byte[read];
                System.arraycopy(inputBuffer, 0, buffer, 0, read);

                buf = cryptoService.decrypt(CipherAlgorithm.find(helper.getBaseCipherName(contentEncryptionAlgorithm)),
                        CipherParameters.builder().encryptionMethod(helper.getEncryptionMethod(contentEncryptionAlgorithm)).key(secretKey).iv(iv).build(), buffer);
                if (buf != null) {
                    maxBuf = buf.length;
                }
            } catch (CryptoException | NoSuchAlgorithmException | CMSException e) {
                throw new CipherIOException("Error processing stream ", e);
            }
        }
        return maxBuf;
    }

    private byte[] finaliseCipher() throws InvalidCipherTextIOException {
        if (!finalized) {
            finalized = true;
        }
        return null;
    }

    /**
     * Reads data from the underlying stream and processes it with the cipher until the cipher outputs data, and returns
     * the next available byte.
     * 
     * <p>
     * If the underlying stream is exhausted by this call, the cipher will be finalised.
     * </p>
     * 
     * @return the next available byte
     * @throws IOException if there was an error closing the input stream.
     * @throws InvalidCipherTextIOException if the data read from the stream was invalid ciphertext (e.g. the cipher is
     *             an AEAD cipher and the ciphertext tag check fails).
     */
    public int read() throws IOException {
        if (bufOff >= maxBuf) {
            if (nextChunk() < 0) {
                return -1;
            }
        }

        return buf[bufOff++] & 0xff;
    }

    /**
     * Reads data from the underlying stream and processes it with the cipher until the cipher outputs data, and then
     * returns up to <code>len</code> bytes in the provided array.
     * <p>
     * If the underlying stream is exhausted by this call, the cipher will be finalised.
     * </p>
     * 
     * @param b the buffer into which the data is read.
     * @param off the start offset in the destination array <code>b</code>
     * @param len the maximum number of bytes read.
     * @return the total number of bytes read into the buffer, or <code>-1</code> if there is no more data because the
     *         end of the stream has been reached.
     * @throws IOException if there was an error closing the input stream.
     * @throws InvalidCipherTextIOException if the data read from the stream was invalid ciphertext (e.g. the cipher is
     *             an AEAD cipher and the ciphertext tag check fails).
     */
    public int read(byte[] b, int off, int len) throws IOException {
        if (bufOff >= maxBuf) {
            if (nextChunk() < 0) {
                return -1;
            }
        }

        int toSupply = Math.min(len, available());
        System.arraycopy(buf, bufOff, b, off, toSupply);
        bufOff += toSupply;
        return toSupply;
    }

    /**
     * Skip.
     * 
     * @param n the n
     * @return skip value
     * @throws IOException {@link IOException}
     */
    public long skip(long n) throws IOException {
        if (n <= 0) {
            return 0;
        }

        int skip = (int) Math.min(n, available());
        bufOff += skip;
        return skip;
    }

    public int available() throws IOException {
        return maxBuf - bufOff;
    }

    /**
     * Closes the underlying input stream, and then finalises the processing of the data by the cipher.
     *
     * @throws IOException if there was an error closing the input stream.
     * @throws InvalidCipherTextIOException if the data read from the stream was invalid ciphertext (e.g. the cipher is
     *             an AEAD cipher and the ciphertext tag check fails).
     */
    public void close() throws IOException {
        try {
            in.close();
        } finally {
            if (!finalized) {
                // Reset the cipher, discarding any data buffered in it
                // Errors in cipher finalisation trump I/O error closing input
                finaliseCipher();
            }
        }
        maxBuf = bufOff = 0;
    }

    public void mark(int readlimit) {
        
    }

    public void reset() throws IOException {
        
    }

    public boolean markSupported() {
        return false;
    }
}
